package com.mm.cloth.rmbg;

import ai.onnxruntime.*;
import android.content.Context;
import android.content.res.AssetManager;
import android.graphics.Bitmap;
import android.graphics.Color;
import android.util.Log;

import java.io.InputStream;
import java.nio.FloatBuffer;
import java.util.Collections;

/**
 * 基于U2Net的衣服分割推理助手
 * 使用ONNX Runtime进行衣服分割推理
 */
public class ClothSegmentationHelper {
    private static final String TAG = "ClothSegmentationHelper";
    
    private OrtSession session;
    private OrtEnvironment env;
    
    // 模型参数
    private static final String INPUT_NAME = "input";
    private static final String OUTPUT_NAME = "output";
    
    // U2Net模型输入尺寸
    private static final int MODEL_INPUT_SIZE = 768;
    
    // ImageNet标准化参数
    private static final float[] MEAN = {0.485f, 0.456f, 0.406f};
    private static final float[] STD = {0.229f, 0.224f, 0.225f};
    
    /**
     * 初始化ONNX Runtime环境
     * @param context 应用上下文
     * @param modelPath 模型文件路径
     */
    public void initModel(Context context, String modelPath) throws Exception {
        Log.d(TAG, "初始化U2Net衣服分割模型，路径: " + modelPath);
        
        try {
            // 检查ONNX Runtime是否可用
            try {
                Class.forName("ai.onnxruntime.OrtEnvironment");
                Log.d(TAG, "ONNX Runtime类可用");
            } catch (ClassNotFoundException e) {
                throw new Exception("ONNX Runtime不可用", e);
            }
            
            // 创建ONNX Runtime环境
            env = OrtEnvironment.getEnvironment();
            Log.d(TAG, "ONNX Runtime环境创建成功");
            
            // 配置会话选项（根据参考文档优化）
            OrtSession.SessionOptions options = new OrtSession.SessionOptions();
            try {
                // 启用所有优化
                options.setOptimizationLevel(OrtSession.SessionOptions.OptLevel.ALL_OPT);
                options.setIntraOpNumThreads(4); // 设置线程数
                
                // 尝试启用NPU加速（华为设备）
                try {
                    options.addConfigEntry("session.use_npu", "1");
                    Log.d(TAG, "尝试启用NPU加速");
                } catch (Exception e) {
                    Log.d(TAG, "NPU不可用，使用CPU");
                }
                
                Log.d(TAG, "会话选项配置完成");
            } catch (Exception e) {
                Log.e(TAG, "会话选项配置失败", e);
                throw new Exception("会话选项配置失败: " + e.getMessage(), e);
            }
            
            // 加载模型文件
            InputStream modelStream = context.getAssets().open(modelPath);
            int available = modelStream.available();
            Log.d(TAG, "模型文件大小: " + (available / (1024*1024)) + " MB");
            
            if (available == 0) {
                modelStream.close();
                throw new Exception("模型文件为空");
            }
            
            // 加载模型数据
            byte[] modelBytes = new byte[available];
            int totalRead = 0;
            while (totalRead < available) {
                int read = modelStream.read(modelBytes, totalRead, available - totalRead);
                if (read == -1) break;
                totalRead += read;
            }
            modelStream.close();
            
            Log.d(TAG, "模型数据加载完成: " + totalRead + " bytes");
            
            // 创建会话
            try {
                Log.d(TAG, "开始创建ONNX会话...");
                session = env.createSession(modelBytes, options);
                Log.d(TAG, "ONNX会话创建成功");
                
                if (session == null) {
                    throw new Exception("会话创建返回null");
                }
                
            } catch (Exception e) {
                Log.e(TAG, "ONNX会话创建失败", e);
                throw new Exception("ONNX会话创建失败: " + e.getMessage(), e);
            }
            
            Log.d(TAG, "U2Net衣服分割模型加载成功");
            
            // 打印模型信息
            printModelInfo();
            
        } catch (Exception e) {
            Log.e(TAG, "U2Net衣服分割模型初始化失败", e);
            // 清理资源
            if (session != null) {
                try { session.close(); } catch (Exception ignored) {}
                session = null;
            }
            if (env != null) {
                try { env.close(); } catch (Exception ignored) {}
                env = null;
            }
            throw e;
        }
    }
    
    /**
     * 打印模型信息
     */
    private void printModelInfo() {
        try {
            Log.d(TAG, "=== U2Net衣服分割模型信息 ===");
            Log.d(TAG, "输入节点数量: " + session.getInputNames().size());
            Log.d(TAG, "输出节点数量: " + session.getOutputNames().size());
            
            for (String inputName : session.getInputNames()) {
                Log.d(TAG, "输入节点: " + inputName);
            }
            
            for (String outputName : session.getOutputNames()) {
                Log.d(TAG, "输出节点: " + outputName);
            }
            
        } catch (Exception e) {
            Log.w(TAG, "获取模型信息失败", e);
        }
    }
    
    /**
     * 执行衣服分割推理
     * @param inputBitmap 输入图像
     * @return 分割掩膜
     */
    public Bitmap segmentCloth(Bitmap inputBitmap) throws Exception {
        if (session == null || env == null) {
            throw new IllegalStateException("模型未初始化");
        }
        
        if (inputBitmap == null) {
            throw new IllegalArgumentException("输入图像不能为空");
        }
        
        Log.d(TAG, "开始U2Net衣服分割推理");
        Log.d(TAG, "输入图像尺寸: " + inputBitmap.getWidth() + "x" + inputBitmap.getHeight());
        
        long startTime = System.currentTimeMillis();
        OnnxTensor tensor = null;
        
        try {
            // 1. 预处理图像
            float[] inputData = preprocessImage(inputBitmap);
            long[] shape = {1, 3, MODEL_INPUT_SIZE, MODEL_INPUT_SIZE}; // NCHW格式
            
            Log.d(TAG, "预处理完成，输入数据长度: " + inputData.length);
            
            // 2. 构建输入张量
            try {
                Log.d(TAG, "创建输入张量...");
                tensor = OnnxTensor.createTensor(env, FloatBuffer.wrap(inputData), shape);
                Log.d(TAG, "输入张量创建成功");
            } catch (Exception e) {
                Log.e(TAG, "输入张量创建失败", e);
                throw new Exception("输入张量创建失败: " + e.getMessage(), e);
            }
            
            // 3. 执行推理
            long inferenceStart = System.currentTimeMillis();
            Log.d(TAG, "开始ONNX推理...");
            try (OrtSession.Result result = session.run(Collections.singletonMap(INPUT_NAME, tensor))) {
                long inferenceEnd = System.currentTimeMillis();
                Log.d(TAG, "ONNX推理完成，耗时: " + (inferenceEnd - inferenceStart) + "ms");
                
                // 4. 获取输出结果
                try {
                    Log.d(TAG, "获取推理结果...");
                    Object outputObject = result.get(0).getValue();
                    if (outputObject == null) {
                        throw new Exception("推理结果为null");
                    }
                    
                    Log.d(TAG, "转换输出数据...");
                    float[] outputData = convertOutputToFloatArray(outputObject);
                    if (outputData == null || outputData.length == 0) {
                        throw new Exception("输出数据为空");
                    }
                    
                    // 5. 后处理：将输出转换为掩膜
                    Log.d(TAG, "开始后处理...");
                    Bitmap maskBitmap = postprocessMask(outputData, inputBitmap.getWidth(), inputBitmap.getHeight());
                    if (maskBitmap == null) {
                        throw new Exception("后处理失败，结果为null");
                    }
                    
                    long totalTime = System.currentTimeMillis() - startTime;
                    Log.d(TAG, "U2Net衣服分割推理完成，总耗时: " + totalTime + "ms");
                    
                    return maskBitmap;
                    
                } catch (Exception e) {
                    Log.e(TAG, "输出处理失败", e);
                    throw new Exception("输出处理失败: " + e.getMessage(), e);
                }
            }
            
        } catch (OutOfMemoryError e) {
            Log.e(TAG, "U2Net推理内存不足", e);
            System.gc();
            throw new Exception("内存不足，请尝试使用较小的图片", e);
        } catch (Exception e) {
            Log.e(TAG, "U2Net推理失败", e);
            throw e;
        } finally {
            // 确保释放张量资源
            if (tensor != null) {
                try {
                    tensor.close();
                } catch (Exception e) {
                    Log.w(TAG, "释放张量资源失败", e);
                }
            }
        }
    }
    
    /**
     * 预处理图像
     * 转换为CHW格式并进行标准化
     */
    private float[] preprocessImage(Bitmap bitmap) {
        // 调整图像尺寸到模型输入尺寸
        Bitmap resizedBitmap = Bitmap.createScaledBitmap(bitmap, MODEL_INPUT_SIZE, MODEL_INPUT_SIZE, true);
        
        // 获取像素数据
        int[] pixels = new int[MODEL_INPUT_SIZE * MODEL_INPUT_SIZE];
        resizedBitmap.getPixels(pixels, 0, MODEL_INPUT_SIZE, 0, 0, MODEL_INPUT_SIZE, MODEL_INPUT_SIZE);
        
        // 转换为CHW格式的float数组
        int totalPixels = MODEL_INPUT_SIZE * MODEL_INPUT_SIZE;
        float[] inputData = new float[3 * totalPixels];
        
        // 分别处理R、G、B三个通道
        for (int c = 0; c < 3; c++) {
            for (int h = 0; h < MODEL_INPUT_SIZE; h++) {
                for (int w = 0; w < MODEL_INPUT_SIZE; w++) {
                    int pixelIndex = h * MODEL_INPUT_SIZE + w;
                    int pixel = pixels[pixelIndex];
                    
                    float value;
                    if (c == 0) { // R通道
                        value = Color.red(pixel) / 255.0f;
                    } else if (c == 1) { // G通道
                        value = Color.green(pixel) / 255.0f;
                    } else { // B通道
                        value = Color.blue(pixel) / 255.0f;
                    }
                    
                    // ImageNet标准化: (value - mean) / std
                    value = (value - MEAN[c]) / STD[c];
                    
                    // CHW格式存储
                    int outputIndex = c * totalPixels + pixelIndex;
                    inputData[outputIndex] = value;
                }
            }
        }
        
        // 清理临时bitmap
        if (resizedBitmap != bitmap) {
            resizedBitmap.recycle();
        }
        
        // 统计输入数据范围，用于调试
        float minVal = Float.MAX_VALUE, maxVal = Float.MIN_VALUE;
        for (float val : inputData) {
            minVal = Math.min(minVal, val);
            maxVal = Math.max(maxVal, val);
        }
        Log.d(TAG, "输入数据范围: [" + minVal + ", " + maxVal + "]");
        
        return inputData;
    }
    
    /**
     * 转换ONNX输出为float数组
     */
    private float[] convertOutputToFloatArray(Object outputObject) {
        try {
            if (outputObject instanceof float[][][][]) {
                // 4维张量格式 [batch, channel, height, width]
                float[][][][] output4D = (float[][][][]) outputObject;
                
                int batch = output4D.length;
                int channels = output4D[0].length;
                int height = output4D[0][0].length;
                int width = output4D[0][0][0].length;
                
                Log.d(TAG, "输出张量形状: [" + batch + ", " + channels + ", " + height + ", " + width + "]");
                
                // U2Net通常输出单通道掩膜
                float[] result = new float[height * width];
                for (int h = 0; h < height; h++) {
                    for (int w = 0; w < width; w++) {
                        result[h * width + w] = output4D[0][0][h][w]; // 取第一个通道
                    }
                }
                
                return result;
                
            } else if (outputObject instanceof float[]) {
                return (float[]) outputObject;
            } else {
                throw new RuntimeException("不支持的输出格式: " + outputObject.getClass().getName());
            }
            
        } catch (Exception e) {
            Log.e(TAG, "输出转换失败", e);
            throw new RuntimeException("输出转换失败", e);
        }
    }

    /**
     * 后处理掩膜数据
     * @param maskData 模型输出的掩膜数据
     * @param targetWidth 目标图像宽度
     * @param targetHeight 目标图像高度
     * @return 处理后的掩膜Bitmap
     */
    private Bitmap postprocessMask(float[] maskData, int targetWidth, int targetHeight) {
        // 首先创建模型输出尺寸的掩膜
        Bitmap modelMask = Bitmap.createBitmap(MODEL_INPUT_SIZE, MODEL_INPUT_SIZE, Bitmap.Config.ARGB_8888);
        int[] pixels = new int[MODEL_INPUT_SIZE * MODEL_INPUT_SIZE];

        // 统计掩膜值范围
        float minVal = Float.MAX_VALUE, maxVal = Float.MIN_VALUE;
        for (float val : maskData) {
            minVal = Math.min(minVal, val);
            maxVal = Math.max(maxVal, val);
        }
        Log.d(TAG, "掩膜值范围: [" + minVal + ", " + maxVal + "]");

        // 将掩膜数据转换为像素
        for (int i = 0; i < maskData.length; i++) {
            float maskValue = maskData[i];

            // 应用sigmoid激活函数（如果输出没有经过sigmoid）
            if (maxVal > 1.0f || minVal < 0.0f) {
                maskValue = 1.0f / (1.0f + (float) Math.exp(-maskValue));
            }

            // 确保值在[0,1]范围内
            maskValue = Math.max(0.0f, Math.min(1.0f, maskValue));

            // 转换为灰度值
            int grayValue = (int) (maskValue * 255);

            // 创建ARGB像素（白色表示衣服，黑色表示背景）
            pixels[i] = Color.argb(255, grayValue, grayValue, grayValue);
        }

        modelMask.setPixels(pixels, 0, MODEL_INPUT_SIZE, 0, 0, MODEL_INPUT_SIZE, MODEL_INPUT_SIZE);

        // 调整到目标尺寸
        Bitmap resultMask;
        if (targetWidth != MODEL_INPUT_SIZE || targetHeight != MODEL_INPUT_SIZE) {
            resultMask = Bitmap.createScaledBitmap(modelMask, targetWidth, targetHeight, true);
            modelMask.recycle();
        } else {
            resultMask = modelMask;
        }

        Log.d(TAG, "掩膜后处理完成，生成掩膜: " + resultMask.getWidth() + "x" + resultMask.getHeight());
        return resultMask;
    }

    /**
     * 应用掩膜到原图，提取衣服区域
     * @param originalImage 原始图像
     * @param maskImage 分割掩膜
     * @return 衣服分割结果
     */
    public Bitmap applyMaskToImage(Bitmap originalImage, Bitmap maskImage) {
        if (originalImage == null || maskImage == null) {
            Log.e(TAG, "原始图像或掩膜为空");
            return null;
        }

        try {
            // 创建结果图像
            Bitmap result = Bitmap.createBitmap(originalImage.getWidth(), originalImage.getHeight(), Bitmap.Config.ARGB_8888);

            // 获取像素数据
            int width = originalImage.getWidth();
            int height = originalImage.getHeight();
            int[] originalPixels = new int[width * height];
            int[] maskPixels = new int[width * height];
            int[] resultPixels = new int[width * height];

            originalImage.getPixels(originalPixels, 0, width, 0, 0, width, height);

            // 调整掩膜尺寸（如果需要）
            Bitmap resizedMask = maskImage;
            if (maskImage.getWidth() != width || maskImage.getHeight() != height) {
                resizedMask = Bitmap.createScaledBitmap(maskImage, width, height, true);
            }
            resizedMask.getPixels(maskPixels, 0, width, 0, 0, width, height);

            // 应用掩膜
            for (int i = 0; i < originalPixels.length; i++) {
                int originalPixel = originalPixels[i];
                int maskPixel = maskPixels[i];

                // 获取掩膜的亮度（白色=衣服，黑色=背景）
                int maskBrightness = Color.red(maskPixel); // 使用红色通道作为亮度

                if (maskBrightness > 128) {
                    // 白色区域：保留原像素（衣服区域）
                    resultPixels[i] = originalPixel;
                } else {
                    // 黑色区域：设置为透明（背景区域）
                    resultPixels[i] = Color.TRANSPARENT;
                }
            }

            // 设置结果像素
            result.setPixels(resultPixels, 0, width, 0, 0, width, height);

            // 清理临时资源
            if (resizedMask != maskImage) {
                resizedMask.recycle();
            }

            Log.d(TAG, "掩膜应用完成，结果尺寸: " + result.getWidth() + "x" + result.getHeight());
            return result;

        } catch (Exception e) {
            Log.e(TAG, "应用掩膜失败", e);
            return null;
        }
    }

    /**
     * 释放资源
     */
    public void release() {
        try {
            if (session != null) {
                session.close();
                session = null;
            }
            if (env != null) {
                env.close();
                env = null;
            }
            Log.d(TAG, "U2Net衣服分割ONNX Runtime资源已释放");
        } catch (Exception e) {
            Log.e(TAG, "释放U2Net衣服分割ONNX Runtime资源失败", e);
        }
    }

    /**
     * 检查是否已初始化
     */
    public boolean isInitialized() {
        return session != null && env != null;
    }
}
